import numpy as np
import random
from sklearn.datasets import make_blobs
from matplotlib import pyplot as plt

X, y = make_blobs(centers=4)
X_list = X.tolist()

# 随机初始化了4个聚类中心
centers_initial = random.sample(population=X_list, k=4)
centers_initial = np.array(centers_initial)
old_classes = []

# 训练过程：目的是为了找到稳定的聚类中心
while True:
    new_classes = []
    plt.scatter(X[:, 0], X[:, 1], c=y)
    plt.scatter(centers_initial[:, 0], centers_initial[:, 1], s=100, marker="*", c="red")
    plt.show()
    for sample in X:
        min_index = np.argmin(((centers_initial - sample) ** 2).sum(axis=1))
        new_classes.append(min_index)
    if new_classes == old_classes:
        print("稳定分类，打印分类中心")
        print(centers_initial)
        break
    else:
        # 重新确定分类中心
        # 重新分类。
        old_classes = new_classes.copy()
        for i in range(4):
            centers_initial[i] = X[np.array(new_classes) == i].mean(axis=0)

# 写出预测过程，实际上就是计算样本到哪个聚类中心近，就属于哪个聚类中心。
def predict(X_test):
    y_predict = []
    for sample in X_test:
        min_index = np.argmin(((centers_initial - sample) ** 2).sum(axis=1))
        y_predict.append(min_index)
    return y_predict
    